#include "compensator.hpp"

#include "opencv2/opencv.hpp"
#include "spdlog/spdlog.h"

namespace {

const double kG = 9.80665;
const cv::Scalar kGREEN(0., 255., 0.);
const cv::Scalar kRED(0., 0., 255.);
const cv::Scalar kYELLOW(0., 255., 255.);
const auto kCV_FONT = cv::FONT_HERSHEY_SIMPLEX;

}  // namespace

/**
 * @brief Angle θ required to hit coordinate (x, y)
 *
 * {\displaystyle \tan \theta ={\left({\frac {v^{2}\pm {\sqrt
 * {v^{4}-g(gx^{2}+2yv^{2})}}}{gx}}\right)}}
 *
 * @param target 目标坐标
 * @return double 出射角度
 */
// double Compensator::SolveSurfaceLanchAngle(cv::Point2f target) {
//   const double v_2 = pow(ballet_speed_, 2);
//   const double up_base =
//       std::sqrt(std::pow(ballet_speed_, 4) -
//                 kG * (kG * std::pow(target.x, 2) + 2 * target.y * v_2));
//   const double low = kG * target.x;
//   const double ans1 = std::atan2(v_2 + up_base, low);
//   const double ans2 = std::atan2(v_2 - up_base, low);

//   if (std::isnan(ans1)) return std::isnan(ans2) ? 0. : ans2;
//   if (std::isnan(ans2)) return std::isnan(ans1) ? 0. : ans1;
//   return std::min(ans1, ans2);
// }

// void Compensator::VisualizePnp(Armor& armor, const cv::Mat& output,
//                                bool add_lable) {
//   std::vector<cv::Point2f> out_points;
//   cv::projectPoints(armor.ImageVertices(), armor.GetRotVec(),
//                     armor.GetTransVec(), cam_mat_, distor_coff_, out_points);
//   for (std::size_t i = 0; i < out_points.size(); ++i) {
//     cv::line(output, out_points[i], out_points[(i + 1) % out_points.size()],
//              kYELLOW);
//   }
//   if (add_lable) {
//     cv::putText(output, "PNP", out_points[0], kCV_FONT, 1.0, kYELLOW);
//   }
// }

Compensator::Compensator() { SPDLOG_TRACE("Constructed."); }

Compensator::Compensator(const std::string& cam_mat_path) {
  SPDLOG_TRACE("Constructed.");
  LoadCameraMat(cam_mat_path);
}

Compensator::~Compensator() { SPDLOG_TRACE("Destructed."); }

void Compensator::LoadCameraMat(const std::string& path) {
  cv::FileStorage fs(path,
                     cv::FileStorage::READ | cv::FileStorage::FORMAT_JSON);

  if (fs.isOpened()) {
    cam_mat_ = fs["cam_mat"].mat();
    distor_coff_ = fs["distor_coff"].mat();
    if (cam_mat_.empty() && distor_coff_.empty()) {
      SPDLOG_ERROR("Can not load cali data.");
    } else {
      SPDLOG_DEBUG("Loaded cali data.");
    }
  } else {
    SPDLOG_ERROR("Can not open file: '{}'", path);
  }
}

// cv::Vec3f Compensator::EstimateWorldCoord(Armor& armor) {
//   cv::Mat rot_vec, trans_vec;
//   cv::solvePnP(armor.PhysicVertices(), armor.ImageVertices(), cam_mat_,
//                distor_coff_, rot_vec, trans_vec, false,
//                cv::SOLVEPNP_ITERATIVE);
//   armor.SetRotVec(rot_vec), armor.SetTransVec(trans_vec);
//   cv::Mat world_coord =
//       ((cv::Vec2f(armor.ImageCenter()) * cam_mat_.inv() - trans_vec) *
//        armor.GetRotMat().inv());
//   return cv::Vec3f(world_coord);
// }

void Compensator::SolveAngles(Armor& armor) {
  cv::Mat rot_vec, trans_vec;
  cv::solvePnP(armor.PhysicVertices(), armor.ImageVertices(), cam_mat_,
               distor_coff_, rot_vec, trans_vec, false, cv::SOLVEPNP_ITERATIVE);
  gun_cam_distance = 0;  //枪管和相机的距离 根据兵种设置
  trans_vec.at<double>(1, 0) -=
      gun_cam_distance;  //因为枪管在摄像头上方 所以ty减去该距离
  armor.SetRotVec(rot_vec), armor.SetTransVec(trans_vec);

  double x_pos = armor.GetTransVec().at<double>(0, 0);
  double y_pos = armor.GetTransVec().at<double>(1, 0);
  double z_pos = armor.GetTransVec().at<double>(2, 0);
  distance = sqrt(x_pos * x_pos + y_pos * y_pos + z_pos * z_pos);

  if (distance > 5000) {
    PinHoleSolver(armor);
  } else {
    P4PSolver(armor);
  }
}

void Compensator::P4PSolver(Armor& armor) {
  double x_pos = armor.GetTransVec().at<double>(0, 0);
  double y_pos = armor.GetTransVec().at<double>(1, 0);
  double z_pos = armor.GetTransVec().at<double>(2, 0);

  double tan_pitch = y_pos / sqrt(x_pos * x_pos + z_pos * z_pos);
  double tan_yaw = x_pos / z_pos;
  component::Euler aiming_eulr;
  aiming_eulr.pitch = -atan(tan_pitch) / CV_PI * 180;
  aiming_eulr.yaw = -atan(tan_yaw) / CV_PI * 180;
  armor.SetAimEuler(aiming_eulr);
}

void Compensator::PinHoleSolver(Armor& armor) {
  double ax = cam_mat_.at<double>(0, 0);
  double ay = cam_mat_.at<double>(1, 1);
  double u0 = cam_mat_.at<double>(0, 2);
  double v0 = cam_mat_.at<double>(1, 2);

  cv::Point2f pnt;
  std::vector<cv::Point2f> in;
  std::vector<cv::Point2f> out;
  in.push_back(armor.image_center_);

  cv::undistortPoints(in, out, cam_mat_, distor_coff_, cv::noArray(), cam_mat_);
  pnt = out.front();

  double tan_pitch = (pnt.y - v0) / ay;
  double tan_yaw = (pnt.x - u0) / ax;
  component::Euler aiming_eulr;
  aiming_eulr.pitch = atan(tan_pitch) / CV_PI * 180;
  aiming_eulr.yaw = -atan(tan_yaw) / CV_PI * 180;
  armor.SetAimEuler(aiming_eulr);
}

void Compensator::Apply(std::vector<Armor>& armors, const cv::Mat& frame) {
  for (auto& armor : armors) {
    // const auto cam_coord = EstimateWorldCoord(armor);
    // const cv::Point3f real_coord(cam_coord.dot(rot_mat));

    // component::Euler aiming_eulr;
    // cv::Point2f surface_target(
    //     std::sqrt(std::pow(real_coord.x, 2) + std::pow(real_coord.y, 2)),
    //     -real_coord.z);
    // aiming_eulr.pitch = SolveSurfaceLanchAngle(surface_target);
    // aiming_eulr.yaw = std::atan2(real_coord.x, real_coord.y);
    // armor.SetAimEuler(aiming_eulr);
    SolveAngles(armor);
    CompensateGravity(armor);
  }
  cv::Point2f frame_center(frame.cols / 2, frame.rows / 2);
  std::sort(armors.begin(), armors.end(),
            [frame_center](Armor& armor1, Armor& armor2) {
              return cv::norm(armor1.ImageCenter() - frame_center) <
                     cv::norm(armor2.ImageCenter() - frame_center);
            });
}

// void Compensator::VisualizeResult(std::vector<Armor>& armors,
//                                   const cv::Mat& output, int verbose) {
//   for (auto& armor : armors) {
//     VisualizePnp(armor, output, verbose > 1);
//   }
// }

void Compensator::CompensateGravity(Armor& armor) {
  component::Euler aiming_eulr;
  aiming_eulr.pitch = armor.GetAimEuler().pitch;
  float compensate_pitch = atan((sin(aiming_eulr.pitch) -
                                 0.5 * kG * distance / pow(ballet_speed_, 2)) /
                                cos(aiming_eulr.pitch)) /
                           CV_PI * 180;
  aiming_eulr.pitch = aiming_eulr.pitch + compensate_pitch;
  armor.SetAimEuler(aiming_eulr);
}
